//! 内存管理模块
//! 
//! 提供高效的内存池管理和零拷贝技术

use std::collections::HashMap;
use std::sync::{Arc, Mutex, RwLock};
use std::time::{Duration, Instant};

/// 内存池配置
#[derive(Debug, Clone)]
pub struct MemoryPoolConfig {
    /// 初始池大小
    pub initial_size: usize,
    /// 最大池大小
    pub max_size: usize,
    /// 缓冲区大小分类
    pub size_classes: Vec<usize>,
    /// 清理间隔
    pub cleanup_interval: Duration,
    /// 最大空闲时间
    pub max_idle_time: Duration,
}

impl Default for MemoryPoolConfig {
    fn default() -> Self {
        Self {
            initial_size: 10,
            max_size: 1000,
            size_classes: vec![64, 256, 1024, 4096, 16384, 65536],
            cleanup_interval: Duration::from_secs(60),
            max_idle_time: Duration::from_secs(300),
        }
    }
}

/// 缓冲区元数据
#[derive(Debug)]
struct BufferMetadata {
    /// 创建时间
    created_at: Instant,
    /// 最后使用时间
    last_used: Instant,
    /// 使用次数
    use_count: usize,
}

/// 内存池
pub struct MemoryPool {
    /// 按大小分类的缓冲区池
    pools: RwLock<HashMap<usize, Vec<(Arc<[u8]>, BufferMetadata)>>>,
    /// 配置
    config: MemoryPoolConfig,
    /// 统计信息
    stats: Arc<Mutex<MemoryStats>>,
    /// 最后清理时间
    last_cleanup: Mutex<Instant>,
}

/// 内存统计信息
#[derive(Debug, Default)]
pub struct MemoryStats {
    /// 总分配次数
    pub total_allocations: usize,
    /// 总释放次数
    pub total_deallocations: usize,
    /// 缓存命中次数
    pub cache_hits: usize,
    /// 缓存未命中次数
    pub cache_misses: usize,
    /// 当前活跃缓冲区数量
    pub active_buffers: usize,
    /// 池中缓冲区数量
    pub pooled_buffers: usize,
    /// 总内存使用量
    pub total_memory_bytes: usize,
}

impl MemoryPool {
    /// 创建新的内存池
    pub fn new(config: MemoryPoolConfig) -> Self {
        let mut pools = HashMap::new();
        
        // 为每个大小类别预分配缓冲区
        for &size in &config.size_classes {
            let mut buffers = Vec::with_capacity(config.initial_size);
            for _ in 0..config.initial_size {
                let buffer = Arc::from(vec![0u8; size].into_boxed_slice());
                let metadata = BufferMetadata {
                    created_at: Instant::now(),
                    last_used: Instant::now(),
                    use_count: 0,
                };
                buffers.push((buffer, metadata));
            }
            pools.insert(size, buffers);
        }

        Self {
            pools: RwLock::new(pools),
            config,
            stats: Arc::new(Mutex::new(MemoryStats::default())),
            last_cleanup: Mutex::new(Instant::now()),
        }
    }

    /// 分配指定大小的缓冲区
    pub fn allocate(&self, size: usize) -> Arc<[u8]> {
        let size_class = self.find_size_class(size);
        let mut stats = self.stats.lock().unwrap();
        stats.total_allocations += 1;

        // 尝试从池中获取缓冲区
        if let Some(buffer) = self.try_get_from_pool(size_class) {
            stats.cache_hits += 1;
            stats.active_buffers += 1;
            return buffer;
        }

        // 池中没有可用缓冲区，创建新的
        stats.cache_misses += 1;
        stats.active_buffers += 1;
        stats.total_memory_bytes += size_class;
        
        Arc::from(vec![0u8; size_class].into_boxed_slice())
    }

    /// 释放缓冲区
    pub fn deallocate(&self, buffer: Arc<[u8]>) {
        let size = buffer.len();
        let mut stats = self.stats.lock().unwrap();
        stats.total_deallocations += 1;
        stats.active_buffers = stats.active_buffers.saturating_sub(1);

        // 检查是否应该放回池中
        if self.should_pool_buffer(size) {
            self.return_to_pool(buffer, size);
            stats.pooled_buffers += 1;
        } else {
            stats.total_memory_bytes = stats.total_memory_bytes.saturating_sub(size);
        }

        // 定期清理
        self.maybe_cleanup();
    }

    /// 查找合适的大小类别
    fn find_size_class(&self, size: usize) -> usize {
        self.config.size_classes
            .iter()
            .find(|&&class_size| class_size >= size)
            .copied()
            .unwrap_or_else(|| {
                // 如果没有合适的大小类别，使用下一个2的幂
                let mut class_size = 1;
                while class_size < size {
                    class_size <<= 1;
                }
                class_size
            })
    }

    /// 尝试从池中获取缓冲区
    fn try_get_from_pool(&self, size_class: usize) -> Option<Arc<[u8]>> {
        let mut pools = self.pools.write().unwrap();
        
        if let Some(buffers) = pools.get_mut(&size_class) {
            if let Some((buffer, mut metadata)) = buffers.pop() {
                metadata.last_used = Instant::now();
                metadata.use_count += 1;
                return Some(buffer);
            }
        }
        
        None
    }

    /// 将缓冲区返回到池中
    fn return_to_pool(&self, buffer: Arc<[u8]>, size: usize) {
        let size_class = self.find_size_class(size);
        let mut pools = self.pools.write().unwrap();
        
        let buffers = pools.entry(size_class).or_insert_with(Vec::new);
        
        // 检查池是否已满
        if buffers.len() < self.config.max_size {
            let metadata = BufferMetadata {
                created_at: Instant::now(),
                last_used: Instant::now(),
                use_count: 0,
            };
            buffers.push((buffer, metadata));
        }
    }

    /// 判断是否应该将缓冲区放回池中
    fn should_pool_buffer(&self, size: usize) -> bool {
        // 只有在配置的大小类别范围内的缓冲区才会被池化
        self.config.size_classes.contains(&self.find_size_class(size))
    }

    /// 可能执行清理操作
    fn maybe_cleanup(&self) {
        let mut last_cleanup = self.last_cleanup.lock().unwrap();
        let now = Instant::now();
        
        if now.duration_since(*last_cleanup) >= self.config.cleanup_interval {
            *last_cleanup = now;
            drop(last_cleanup); // 释放锁
            self.cleanup_expired_buffers();
        }
    }

    /// 清理过期的缓冲区
    fn cleanup_expired_buffers(&self) {
        let mut pools = self.pools.write().unwrap();
        let now = Instant::now();
        let mut stats = self.stats.lock().unwrap();
        
        for (size_class, buffers) in pools.iter_mut() {
            let initial_count = buffers.len();
            
            buffers.retain(|(_, metadata)| {
                now.duration_since(metadata.last_used) < self.config.max_idle_time
            });
            
            let removed_count = initial_count - buffers.len();
            stats.pooled_buffers = stats.pooled_buffers.saturating_sub(removed_count);
            stats.total_memory_bytes = stats.total_memory_bytes.saturating_sub(removed_count * size_class);
        }
    }

    /// 获取统计信息
    pub fn stats(&self) -> MemoryStats {
        let stats = self.stats.lock().unwrap();
        stats.clone()
    }

    /// 重置统计信息
    pub fn reset_stats(&self) {
        let mut stats = self.stats.lock().unwrap();
        *stats = MemoryStats::default();
    }

    /// 强制清理所有缓冲区
    pub fn clear(&self) {
        let mut pools = self.pools.write().unwrap();
        let mut stats = self.stats.lock().unwrap();
        
        for buffers in pools.values_mut() {
            buffers.clear();
        }
        
        stats.pooled_buffers = 0;
        stats.total_memory_bytes = 0;
    }

    /// 获取池状态信息
    pub fn pool_info(&self) -> HashMap<usize, usize> {
        let pools = self.pools.read().unwrap();
        pools.iter()
            .map(|(&size, buffers)| (size, buffers.len()))
            .collect()
    }
}

impl Default for MemoryPool {
    fn default() -> Self {
        Self::new(MemoryPoolConfig::default())
    }
}

/// 零拷贝字符串池
pub struct StringPool {
    /// 字符串缓存
    cache: RwLock<HashMap<String, Arc<str>>>,
    /// 最大缓存大小
    max_size: usize,
    /// 访问计数
    access_count: Mutex<HashMap<String, usize>>,
}

impl StringPool {
    /// 创建新的字符串池
    pub fn new(max_size: usize) -> Self {
        Self {
            cache: RwLock::new(HashMap::new()),
            max_size,
            access_count: Mutex::new(HashMap::new()),
        }
    }

    /// 获取字符串，如果不存在则创建
    pub fn get_or_create(&self, s: &str) -> Arc<str> {
        // 首先尝试从缓存中获取
        {
            let cache = self.cache.read().unwrap();
            if let Some(cached) = cache.get(s) {
                // 更新访问计数
                let mut access_count = self.access_count.lock().unwrap();
                *access_count.entry(s.to_string()).or_insert(0) += 1;
                return Arc::clone(cached);
            }
        }

        // 缓存中不存在，需要创建
        let arc_str: Arc<str> = Arc::from(s);
        
        {
            let mut cache = self.cache.write().unwrap();
            
            // 检查缓存大小，如果超过限制则清理
            if cache.len() >= self.max_size {
                self.evict_least_used(&mut cache);
            }
            
            cache.insert(s.to_string(), Arc::clone(&arc_str));
        }

        // 更新访问计数
        {
            let mut access_count = self.access_count.lock().unwrap();
            access_count.insert(s.to_string(), 1);
        }

        arc_str
    }

    /// 清理最少使用的条目
    fn evict_least_used(&self, cache: &mut HashMap<String, Arc<str>>) {
        let access_count = self.access_count.lock().unwrap();
        
        // 找到访问次数最少的条目
        if let Some((least_used_key, _)) = access_count
            .iter()
            .min_by_key(|(_, &count)| count)
        {
            cache.remove(least_used_key);
        }
    }

    /// 获取缓存统计信息
    pub fn cache_info(&self) -> (usize, usize) {
        let cache = self.cache.read().unwrap();
        (cache.len(), self.max_size)
    }

    /// 清空缓存
    pub fn clear(&self) {
        let mut cache = self.cache.write().unwrap();
        let mut access_count = self.access_count.lock().unwrap();
        
        cache.clear();
        access_count.clear();
    }
}

impl Default for StringPool {
    fn default() -> Self {
        Self::new(1000)
    }
}

impl Clone for MemoryStats {
    fn clone(&self) -> Self {
        Self {
            total_allocations: self.total_allocations,
            total_deallocations: self.total_deallocations,
            cache_hits: self.cache_hits,
            cache_misses: self.cache_misses,
            active_buffers: self.active_buffers,
            pooled_buffers: self.pooled_buffers,
            total_memory_bytes: self.total_memory_bytes,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_memory_pool_allocation() {
        let pool = MemoryPool::default();
        
        let buffer = pool.allocate(1024);
        assert_eq!(buffer.len(), 1024);
        
        let stats = pool.stats();
        assert_eq!(stats.total_allocations, 1);
        assert_eq!(stats.active_buffers, 1);
    }

    #[test]
    fn test_memory_pool_deallocation() {
        let pool = MemoryPool::default();
        
        let buffer = pool.allocate(1024);
        pool.deallocate(buffer);
        
        let stats = pool.stats();
        assert_eq!(stats.total_deallocations, 1);
        assert_eq!(stats.active_buffers, 0);
    }

    #[test]
    fn test_memory_pool_reuse() {
        // 创建一个没有预分配缓冲区的内存池
        let config = MemoryPoolConfig {
            initial_size: 0,  // 不预分配
            max_size: 1000,
            size_classes: vec![64, 256, 1024, 4096, 16384, 65536],
            cleanup_interval: Duration::from_secs(60),
            max_idle_time: Duration::from_secs(300),
        };
        let pool = MemoryPool::new(config);
        
        // 第一次分配应该是 cache miss
        let buffer1 = pool.allocate(1024);
        let stats1 = pool.stats();
        assert_eq!(stats1.cache_hits, 0);
        assert_eq!(stats1.cache_misses, 1);
        
        // 释放缓冲区
        pool.deallocate(buffer1);
        
        // 第二次分配应该是 cache hit（重用释放的缓冲区）
        let _buffer2 = pool.allocate(1024);
        let stats2 = pool.stats();
        assert_eq!(stats2.cache_hits, 1);
        assert_eq!(stats2.cache_misses, 1);
    }

    #[test]
    fn test_string_pool() {
        let pool = StringPool::new(10);
        
        let s1 = pool.get_or_create("hello");
        let s2 = pool.get_or_create("hello");
        
        assert!(Arc::ptr_eq(&s1, &s2));
        
        let (cache_size, max_size) = pool.cache_info();
        assert_eq!(cache_size, 1);
        assert_eq!(max_size, 10);
    }

    #[test]
    fn test_string_pool_eviction() {
        let pool = StringPool::new(2);
        
        let _s1 = pool.get_or_create("hello");
        let _s2 = pool.get_or_create("world");
        let _s3 = pool.get_or_create("rust"); // 这应该触发清理
        
        let (cache_size, _) = pool.cache_info();
        assert_eq!(cache_size, 2); // 缓存大小应该保持在限制内
    }

    #[test]
    fn test_size_class_selection() {
        let pool = MemoryPool::default();
        
        assert_eq!(pool.find_size_class(50), 64);
        assert_eq!(pool.find_size_class(100), 256);
        assert_eq!(pool.find_size_class(1000), 1024);
        assert_eq!(pool.find_size_class(100000), 131072); // 下一个2的幂
    }
} 